﻿/* SCE CONFIDENTIAL
$PSLibId$
 * Copyright (C) 2011 Sony Computer Entertainment Inc.
 * All Rights Reserved.
 */

#include "heap.h"
#include <libdbg.h>
#include <stdlib.h>

// Align value
#define ALIGN(x, a)					(((x) + ((a) - 1)) & ~((a) - 1))

struct HeapBlock
{
	HeapBlock *next;
	int32_t type;
	uintptr_t base;
	uint32_t offset;
	uint32_t size;
};

struct HeapContext
{
	HeapBlock *allocList;		///< Not sorted, LIFO.
	HeapBlock *freeList;		///< Sorted by base address, always fully merged.
};

static HeapContext g_context;

static HeapBlock *heapBlockAlloc()
{
	// NOTE: these blocks could be pooled to improve performance
	return (HeapBlock *)malloc(sizeof(HeapBlock));
}

static void heapBlockFree(HeapBlock *block)
{
	// NOTE: these blocks could be pooled to improve performance
	free(block);
}

static bool heapBlockCanMerge(HeapBlock *a, HeapBlock *b)
{
	return a->type				== b->type
		&& a->base + a->size	== b->base
		&& a->offset + a->size	== b->offset;
}

static void heapInsertFreeBlock(HeapContext *ctx, HeapBlock *block)
{
	// advance through free list until prev < block < curr
	HeapBlock *curr = ctx->freeList;
	HeapBlock *prev = NULL;
	while (curr && curr->base < block->base) {
		prev = curr;
		curr = curr->next;
	}

	// sanity check no overlap with adjacent blocks
	SCE_DBG_ASSERT(!prev || prev->base + prev->size <= block->base);
	SCE_DBG_ASSERT(!curr || block->base + block->size <= curr->base);

	// insert into list
	if (prev) {
		prev->next = block;
	} else {
		ctx->freeList = block;
	}
	block->next = curr;

	// check merge with current
	if (curr && heapBlockCanMerge(block, curr)) {
		block->size += curr->size;
		block->next = curr->next;
		heapBlockFree(curr);
	}

	// check merge with prev
	if (prev  && heapBlockCanMerge(prev, block)) {
		prev->size += block->size;
		prev->next = block->next;
		heapBlockFree(block);
	}
}

static HeapBlock *heapAllocBlock(HeapContext *ctx, int32_t type, uint32_t size, uint32_t alignment)
{
	// find a suitable block in the free list
	HeapBlock *curr = ctx->freeList;
	HeapBlock *prev = NULL;
	while (curr) {
		// check this block can handle alignment and size
		uint32_t const skip = ALIGN(curr->base, alignment) - curr->base;
		if (curr->type == type && skip + size <= curr->size) {
			// allocate any blocks we need now to avoid complicated rollback
			HeapBlock *skipBlock = NULL;
			HeapBlock *unusedBlock = NULL;
			if (skip != 0) {
				skipBlock = heapBlockAlloc();
				if (!skipBlock) {
					return NULL;
				}
			}
			if (skip + size != curr->size) {
				unusedBlock = heapBlockAlloc();
				if (!unusedBlock) {
					if (skipBlock) {
						heapBlockFree(skipBlock);
					}
					return NULL;
				}
			}

			// add block for skipped memory
			if (skip != 0) {
				// link in
				if (prev) {
					prev->next = skipBlock;
				} else {
					ctx->freeList = skipBlock;
				}
				skipBlock->next = curr;

				// set sizes
				skipBlock->type = curr->type;
				skipBlock->base = curr->base;
				skipBlock->offset = curr->offset;
				skipBlock->size = skip;
				curr->base += skip;
				curr->offset += skip;
				curr->size -= skip;

				// update prev
				prev = skipBlock;
			}

			// add block for unused memory
			if (size != curr->size) {
				// link in
				unusedBlock->next = curr->next;
				curr->next = unusedBlock;

				// set sizes
				unusedBlock->type = curr->type;
				unusedBlock->base = curr->base + size;
				unusedBlock->offset = curr->offset + size;
				unusedBlock->size = curr->size - size;
				curr->size = size;
			}

			// unlink from free list
			if (prev) {
				prev->next = curr->next;
			} else {
				ctx->freeList = curr->next;
			}

			// push onto alloc list
			curr->next = ctx->allocList;
			ctx->allocList = curr;
			return curr;
		}

		// advance
		prev = curr;
		curr = curr->next;
	}

	// no block found
	return NULL;
}

static void heapFreeBlock(HeapContext *ctx, uintptr_t base)
{
	// find in the allocate block list
	HeapBlock *curr = ctx->allocList;
	HeapBlock *prev = NULL;
	while (curr && curr->base != base) {
		prev = curr;
		curr = curr->next;
	}

	// early out if not found
	if (!curr)
		return;

	// unlink from allocated list
	if (prev) {
		prev->next = curr->next;
	} else {
		ctx->allocList = curr->next;
	}
	curr->next = NULL;

	// add as free block
	heapInsertFreeBlock(&g_context, curr);
}

void heapInitialize()
{
	g_context.allocList = NULL;
	g_context.freeList = NULL;
}

void heapTerminate()
{
	// should be empty
	SCE_DBG_ASSERT(!g_context.allocList);
}

void heapExtend(int32_t type, void *base, uint32_t size, uint32_t offset)
{
	HeapBlock *block = heapBlockAlloc();
	SCE_DBG_ASSERT(block);
	block->next = NULL;
	block->type = type;
	block->base = (uintptr_t)base;
	block->offset = offset;
	block->size = size;
	heapInsertFreeBlock(&g_context, block);
}

void *heapAlloc(int32_t type, uint32_t size, uint32_t alignment, uint32_t *offset)
{
	// try to allocate
	HeapBlock *block = heapAllocBlock(&g_context, type, size, alignment);

	// early out if failed
	if (!block)
		return NULL;

	// write offset and return base address
	if (offset) {
		*offset = block->offset;
	}
	return (void *)block->base;
}

void heapFree(void *addr)
{
	heapFreeBlock(&g_context, (uintptr_t)addr);
}
